import os
import json
import torchvision
import matplotlib.pyplot as plt
from model_repo import *
import dataproc
import numpy as np


train_loader, val_loader = dataproc.dtcustom.custom_dtset()
test_data_iter = iter(val_loader)
test_image, test_label = test_data_iter.next()

# read class_indict
json_path = './class_indices.json'
assert os.path.exists(
    json_path), "file: '{}' dose not exist.".format(json_path)

json_file = open(json_path, "r")
class_indict = json.load(json_file)


def imshow(img):
    img = img / 2 + 0.5
    npimg = img.numpy()
    plt.imshow(np.transpose(npimg, (1, 2, 0)))
    plt.show()


print(' '.join('%5s' % class_indict[test_label[j].item()] for j in range(4)))
imshow(torchvision.utils.make_grid(test_image))
